import {
  StreamlitComponentBase,
  withStreamlitConnection,
  Streamlit,
} from "streamlit-component-lib";
import React, { ReactNode } from "react";

interface Image {
  src: string,
  title: string,
  is_local?: boolean
}

interface GalleryState {
  hoveredIndex: number;
  imageErrors: Set<number>;
  clickedIndex: number;
}

class StreamlitGallery extends StreamlitComponentBase {
  private imageListElement: HTMLDivElement | null | undefined;

  // Arguments that are passed to the plugin in Python are accessible
  private images: Image[] = this.props.args["images"];

  // State to track hover for each image
  state: GalleryState = {
    hoveredIndex: -1,
    imageErrors: new Set<number>(),
    clickedIndex: -1
  };

  private maxWidth = this.props.args["max_width"] ?? 400;

  private gap = this.props.args["gap"] ?? 10;
  private height = this.props.args["height"] ?? 150;

  private maxCols = this.props.args["max_cols"] ?? 2;
  private maxRows = this.props.args["max_rows"] ?? 2;
  private borderColor = this.props.args["border_color"] ?? "none";
  private clickedBorderColor = this.props.args["clicked_border_color"] ?? "#4da6ff";

  private handleImageError = (index: number) => {
    this.setState((prevState: GalleryState) => ({
      imageErrors: new Set([...Array.from(prevState.imageErrors), index])
    }));
  };

  private handleImageLoad = (index: number) => {
    this.setState((prevState: GalleryState) => {
      const newErrors = new Set(Array.from(prevState.imageErrors));
      newErrors.delete(index);
      return { imageErrors: newErrors };
    });
  };

    public render = (): ReactNode => {
    const numberImages = this.images.length;
    const cols = Math.min(numberImages, this.maxCols);

    return (
      <div
        style={{
          display: 'grid',
          gridTemplateColumns: `repeat(${cols}, 1fr)`,
          gap: `${this.gap}px`,
          maxWidth: this.maxWidth,
          width: '100%',
          margin: '0 auto'
        }}
        ref={(element) => { this.imageListElement = element }}
      >
        {this.images.map((image, index) => {
          const hasError = this.state.imageErrors.has(index);
          
          return (
            <div
              key={image.src}
              onClick={(event) => this.openImage(event, index)}
              onMouseEnter={() => this.setState({ hoveredIndex: index })}
              onMouseLeave={() => this.setState({ hoveredIndex: -1 })}
              style={{
                display: 'flex',
                flexDirection: 'column',
                overflow: 'hidden',
                borderRadius: '4px',
                backgroundColor: this.state.clickedIndex === index ? '#e6f3ff' : 
                               (this.state.hoveredIndex === index ? '#f0f8ff' : 'transparent'),
                transition: 'background-color 0.2s ease',
                cursor: 'pointer',
                border: this.state.clickedIndex === index ? '2px solid #4da6ff' : 
                       (this.borderColor ? `2px solid ${this.borderColor}` : 'none')
              }}
            >
              {hasError ? (
                <div
                  style={{
                    height: `${this.height}px`,
                    width: '100%',
                    display: 'flex',
                    alignItems: 'center',
                    justifyContent: 'center',
                    backgroundColor: '#f8f9fa',
                    color: '#6c757d',
                    fontSize: '12px',
                    textAlign: 'center',
                    padding: '10px'
                  }}
                >
                  {image.is_local ? '本地图片加载失败' : '图片加载失败'}
                </div>
              ) : (
                <img
                  src={image.src}
                  alt={image.title}
                  style={{
                    cursor: "pointer",
                    display: 'block',
                    height: `${this.height}px`,
                    width: '100%',
                    margin: '0 auto',
                    objectFit: 'cover'
                  }}
                  loading="lazy"
                  onError={() => this.handleImageError(index)}
                  onLoad={() => this.handleImageLoad(index)}
                />
              )}
              <p
                style={{
                  cursor: "pointer",
                  margin: "0",
                  padding: "12px",
                  textAlign: "center",
                  fontSize: "14px",
                  fontWeight: "500",
                  color: hasError ? "#dc3545" : "#333",
                  borderRadius: "0 0 4px 4px"
                }}
              >
                {image.title}
              </p>
            </div>
          );
        })}
      </div>
    )
  }

  private openImage = (event: React.MouseEvent<HTMLElement>, index: number) => {
    // 更新点击状态
    this.setState({ clickedIndex: index });
    Streamlit.setComponentValue(index);
  }
}

export default withStreamlitConnection(StreamlitGallery);
